package fun.yyds.controller.sys;

import com.baidubce.qianfan.Qianfan;
import com.baidubce.qianfan.QianfanV2;
import com.baidubce.qianfan.core.StreamIterator;
import com.baidubce.qianfan.model.chat.v2.response.StreamResponseV2;
import fun.yyds.common.ResponseData;
import fun.yyds.entity.message.AiMsg;
import fun.yyds.service.ai.AiService;
import io.swagger.v3.oas.annotations.Operation;
import io.swagger.v3.oas.annotations.tags.Tag;
import org.springframework.web.bind.annotation.*;
import org.springframework.web.servlet.mvc.method.annotation.ResponseBodyEmitter;

import javax.annotation.Resource;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import java.io.IOException;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;

@RestController
@Tag(name = "2.Ai")
@RequestMapping("/ai")
public class AIController {
    @Resource
    private AiService aiService;

    @PostMapping("/all/chat")
    @Operation(summary = "非流式普通用户ai对话")
    public ResponseData getAllChatMsg(@RequestBody AiMsg aiMsg) {
        return ResponseData.success(aiService.getAllChatMsg(aiMsg));
    }

    @PostMapping("/my/chat")
    @Operation(summary = "非流式注册用户ai对话")
    public ResponseData getMyChatMsg(@RequestBody AiMsg aiMsg) {
        return ResponseData.success(aiService.getMyChatMsg(aiMsg));
    }

    @PostMapping("/my/stream/chat")
    @Operation(summary = "流式注册用户ai对话")
    public void getMyStreamChatMsg(@RequestBody String message, HttpServletRequest req, HttpServletResponse resp) {
        aiService.getMyStreamChatMsg(message, req, resp);
    }

    @GetMapping("/stream-chat")
    @Operation(summary = "流式注册用户ai对话")
    public ResponseBodyEmitter streamChat(@RequestParam String message) {
        ResponseBodyEmitter emitter = new ResponseBodyEmitter();
        ExecutorService executor = Executors.newSingleThreadExecutor();

        executor.execute(() -> {
            try {
                QianfanV2 client = new Qianfan("2491620a3bc945258221e04beecc8d45",
                        "3cfdfd72fe0744759c7d5c63e59b457d").v2();

                StreamIterator<StreamResponseV2> responses = client.chatCompletionStream(
                        client.chatCompletion()
                                .model("ernie-3.5-8k")
                                .addMessage("user", message)
                                .temperature(0.7)
                                .build()
                );

                // 流式发送数据
                responses.forEachRemaining(resp -> {
                    String content = resp.getChoices().get(0).getDelta().getContent();
                    try {
                        emitter.send(content + " "); // 添加分隔符
                    } catch (IOException e) {
                        throw new RuntimeException("客户端连接已关闭", e);
                    }
                });

                emitter.complete();
            } catch (Exception e) {
                emitter.completeWithError(e);
            } finally {
                executor.shutdown();
            }
        });

        // 处理客户端断开连接
        emitter.onCompletion(executor::shutdownNow);
        emitter.onTimeout(executor::shutdownNow);

        return emitter;
    }

}
